/*
 * Licensed to the Apache Software Foundation (ASF) under one or more
 * contributor license agreements.  See the NOTICE file distributed with
 * this work for additional information regarding copyright ownership.
 * The ASF licenses this file to You under the Apache License, Version 2.0
 * (the "License"); you may not use this file except in compliance with
 * the License.  You may obtain a copy of the License at
 *
 *    http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */

package org.apache.dolphinscheduler.dao.repository.impl;

import org.apache.dolphinscheduler.common.enums.FailureStrategy;
import org.apache.dolphinscheduler.common.enums.Flag;
import org.apache.dolphinscheduler.common.enums.WorkflowExecutionStatus;
import org.apache.dolphinscheduler.common.model.Server;
import org.apache.dolphinscheduler.common.utils.DateUtils;
import org.apache.dolphinscheduler.common.utils.JSONUtils;
import org.apache.dolphinscheduler.dao.entity.ChargingRule;
import org.apache.dolphinscheduler.dao.entity.ProcessInstance;
import org.apache.dolphinscheduler.dao.entity.Product;
import org.apache.dolphinscheduler.dao.entity.ProductOrder;
import org.apache.dolphinscheduler.dao.entity.ProductOrderItem;
import org.apache.dolphinscheduler.dao.entity.TaskInstance;
import org.apache.dolphinscheduler.dao.mapper.ChargingRuleMapper;
import org.apache.dolphinscheduler.dao.mapper.ProcessInstanceMapper;
import org.apache.dolphinscheduler.dao.mapper.ProductMapper;
import org.apache.dolphinscheduler.dao.mapper.ProductOrderItemMapper;
import org.apache.dolphinscheduler.dao.mapper.ProductOrderMapper;
import org.apache.dolphinscheduler.dao.mapper.TaskInstanceMapper;
import org.apache.dolphinscheduler.dao.repository.BaseDao;
import org.apache.dolphinscheduler.dao.repository.TaskInstanceDao;
import org.apache.dolphinscheduler.plugin.task.api.enums.TaskExecutionStatus;
import org.apache.dolphinscheduler.registry.api.RegistryClient;
import org.apache.dolphinscheduler.registry.api.enums.RegistryNodeType;

import org.apache.commons.lang3.StringUtils;

import java.math.BigDecimal;
import java.util.Date;
import java.util.List;

import lombok.NonNull;
import lombok.extern.slf4j.Slf4j;

import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.stereotype.Repository;

import com.alibaba.fastjson.JSONObject;
import com.baomidou.mybatisplus.core.conditions.query.QueryWrapper;

/**
 * Task Instance DAO implementation
 */
@Repository
@Slf4j
public class TaskInstanceDaoImpl extends BaseDao<TaskInstance, TaskInstanceMapper> implements TaskInstanceDao {

    @Autowired
    private ProcessInstanceMapper processInstanceMapper;

    @Autowired
    private ProductOrderMapper productOrderMapper;
    @Autowired
    private ProductOrderItemMapper productOrderItemMapper;
    @Autowired
    private ProductMapper productMapper;
    @Autowired
    private ChargingRuleMapper chargingRuleMapper;
    @Autowired
    private RegistryClient registryClient;

    public TaskInstanceDaoImpl(@NonNull TaskInstanceMapper taskInstanceMapper) {
        super(taskInstanceMapper);
    }

    @Override
    public boolean upsertTaskInstance(TaskInstance taskInstance) {
        boolean result = false;
        if (taskInstance.getId() != null) {
            result = updateById(taskInstance);
            updateProductOrderItem(taskInstance);
        } else {
            result = insert(taskInstance) > 0;
            saveProductOrderItem(taskInstance);
        }
        return result;
    }

    @Override
    public boolean submitTaskInstanceToDB(TaskInstance taskInstance, ProcessInstance processInstance) {
        WorkflowExecutionStatus processInstanceState = processInstance.getState();
        if (processInstanceState.isFinished() || processInstanceState == WorkflowExecutionStatus.READY_STOP) {
            log.warn("processInstance: {} state was: {}, skip submit this task, taskCode: {}",
                    processInstance.getId(),
                    processInstanceState,
                    taskInstance.getTaskCode());
            return false;
        }
        if (processInstanceState == WorkflowExecutionStatus.READY_PAUSE) {
            taskInstance.setState(TaskExecutionStatus.PAUSE);
        }
        taskInstance.setExecutorId(processInstance.getExecutorId());
        taskInstance.setExecutorName(processInstance.getExecutorName());
        taskInstance.setState(getSubmitTaskState(taskInstance, processInstance));
        if (taskInstance.getSubmitTime() == null) {
            taskInstance.setSubmitTime(new Date());
        }
        if (taskInstance.getFirstSubmitTime() == null) {
            taskInstance.setFirstSubmitTime(taskInstance.getSubmitTime());
        }
        return upsertTaskInstance(taskInstance);
    }

    private TaskExecutionStatus getSubmitTaskState(TaskInstance taskInstance, ProcessInstance processInstance) {
        TaskExecutionStatus state = taskInstance.getState();
        if (state == TaskExecutionStatus.RUNNING_EXECUTION
                || state == TaskExecutionStatus.DELAY_EXECUTION
                || state == TaskExecutionStatus.KILL
                || state == TaskExecutionStatus.DISPATCH) {
            return state;
        }

        if (processInstance.getState() == WorkflowExecutionStatus.READY_PAUSE) {
            state = TaskExecutionStatus.PAUSE;
        } else if (processInstance.getState() == WorkflowExecutionStatus.READY_STOP
                || !checkProcessStrategy(taskInstance, processInstance)) {
            state = TaskExecutionStatus.KILL;
        } else {
            state = TaskExecutionStatus.SUBMITTED_SUCCESS;
        }
        return state;
    }

    private boolean checkProcessStrategy(TaskInstance taskInstance, ProcessInstance processInstance) {
        FailureStrategy failureStrategy = processInstance.getFailureStrategy();
        if (failureStrategy == FailureStrategy.CONTINUE) {
            return true;
        }
        List<TaskInstance> taskInstances =
                this.queryValidTaskListByWorkflowInstanceId(taskInstance.getProcessInstanceId(),
                        taskInstance.getTestFlag());

        for (TaskInstance task : taskInstances) {
            if (task.getState() == TaskExecutionStatus.FAILURE
                    && task.getRetryTimes() >= task.getMaxRetryTimes()) {
                return false;
            }
        }
        return true;
    }

    @Override
    public List<TaskInstance> queryValidTaskListByWorkflowInstanceId(Integer processInstanceId, int testFlag) {
        return mybatisMapper.findValidTaskListByProcessId(processInstanceId, Flag.YES, testFlag);
    }

    @Override
    public TaskInstance queryByWorkflowInstanceIdAndTaskCode(Integer processInstanceId, Long taskCode) {
        return mybatisMapper.queryByInstanceIdAndCode(processInstanceId, taskCode);
    }

    @Override
    public List<TaskInstance> queryPreviousTaskListByWorkflowInstanceId(Integer processInstanceId) {
        ProcessInstance processInstance = processInstanceMapper.selectById(processInstanceId);
        return mybatisMapper.findValidTaskListByProcessId(processInstanceId, Flag.NO,
                processInstance.getTestFlag());
    }

    @Override
    public TaskInstance queryByCacheKey(String cacheKey) {
        if (StringUtils.isEmpty(cacheKey)) {
            return null;
        }
        return mybatisMapper.queryByCacheKey(cacheKey);
    }

    @Override
    public Boolean clearCacheByCacheKey(String cacheKey) {
        try {
            mybatisMapper.clearCacheByCacheKey(cacheKey);
            return true;
        } catch (Exception e) {
            log.error("clear cache by cacheKey failed", e);
            return false;
        }
    }

    @Override
    public void deleteByWorkflowInstanceId(int workflowInstanceId) {
        mybatisMapper.deleteByWorkflowInstanceId(workflowInstanceId);
    }

    @Override
    public List<TaskInstance> queryByWorkflowInstanceId(Integer workflowInstanceId) {
        return mybatisMapper.findByWorkflowInstanceId(workflowInstanceId);
    }

    public void updateProductOrderItem(TaskInstance taskInstance) {
        if (taskInstance.getId() == 0) {
            log.error("update product order item error, id is null");
            return;
        }
        ProductOrderItem productOrderItem =
                productOrderItemMapper.queryProductOrderItemByTaskInstanceId(taskInstance.getId());
        if (productOrderItem == null || productOrderItem.getId() == null) {
            log.info("productOrderItem is null, to save");
            saveProductOrderItem(taskInstance);
            return;
        }

        if(taskInstance.getState().equals(TaskExecutionStatus.DELAY_EXECUTION)){
            log.info("taskInstance state is delay execution");
            return;
        }
        if (!productOrderItem.getTaskStatus().equals(0)) {
            log.info("productOrderItem is successed or failured");
            return;
        }

        log.info("============= update product order item (" + taskInstance.getId() + ") start =============");
        log.info(JSONUtils.toJsonString(taskInstance));
        if (taskInstance.getStartTime() != null) {
            productOrderItem.setStartTime(taskInstance.getStartTime());
        }
        if (taskInstance.getEndTime() != null) {
            productOrderItem.setEndTime(taskInstance.getEndTime());
        }
        String host = taskInstance.getHost();
        if (StringUtils.isNotBlank(host) && host.indexOf(":") != -1) {
            host = host.substring(0, host.indexOf(":"));
        }

        // 获取资源利用率
        if (taskInstance.getState().equals(TaskExecutionStatus.RUNNING_EXECUTION)) {
            List<Server> listServer = registryClient.getServerList(RegistryNodeType.WORKER);
            for (Server server : listServer) {
                if (server.getHost().equals(host)) {
                    String info = server.getResInfo();
                    if (StringUtils.isBlank(info)) {
                        break;
                    }
                    JSONObject jsonObject = JSONObject.parseObject(info);
                    if (jsonObject == null || !jsonObject.containsKey("loadAverage")) {
                        break;
                    }
                    double loadAverage = jsonObject.getDouble("loadAverage");
                    log.info(host + ": " + loadAverage);
                    productOrderItem.setResourceUsage(new BigDecimal(loadAverage));
                    break;
                }
            }
            productOrderItemMapper.updateState(productOrderItem);
            return;
        }

        if (StringUtils.isNotBlank(host)) {
            Product product = productMapper.selectProductByHost(host);
            if (product != null) {
                productOrderItem.setResourceNo(product.getProductNo()); // 根据hosts获取资源编号
                productOrderItem.setProductName(product.getProductName());
                productOrderItem.setHost(product.getHost());
                productOrderItem.setProductType(product.getProductType());
                productOrderItem.setProductSource(product.getProductSource());
                productOrderItem.setBillType(product.getBillType());
                productOrderItem.setBillRule(product.getBillRule());
                productOrderItem.setPrice(product.getPrice());
                productOrderItem.setvCpu(product.getvCpu());
                productOrderItem.setMemory(product.getMemory());
                productOrderItem.setDataTotal(product.getDataTotal());
                productOrderItem.setServiceRatio(product.getServiceRatio());
            }
        }
        if (StringUtils.isNotBlank(taskInstance.getHost())
                && taskInstance.getStartTime() != null
                && taskInstance.getEndTime() != null) {
            calcAmount(taskInstance, productOrderItem);
        }
        if (taskInstance.getState().equals(TaskExecutionStatus.SUCCESS)
                || taskInstance.getState().equals(TaskExecutionStatus.FORCED_SUCCESS) ) {
            productOrderItem.setTaskStatus(1);
        } else if(taskInstance.getState().equals(TaskExecutionStatus.FAILURE)
                || taskInstance.getState().equals(TaskExecutionStatus.KILL)
                || taskInstance.getState().equals(TaskExecutionStatus.PAUSE)
                || taskInstance.getState().equals(TaskExecutionStatus.STOP)) {
            productOrderItem.setTaskStatus(2);
        }
        productOrderItem.setStartTime(taskInstance.getStartTime());
        productOrderItem.setEndTime(taskInstance.getEndTime());
        productOrderItemMapper.updateState(productOrderItem);
        log.info("============= update product order item (" + taskInstance.getId() + ") success =============");
    }

    private void saveProductOrderItem(TaskInstance taskInstance) {
        if (taskInstance.getId() == 0) {
            log.error("create product order item error, id is null");
            return;
        }
        ProductOrderItem productOrderItem =
                productOrderItemMapper.queryProductOrderItemByTaskInstanceId(taskInstance.getId());
        if (productOrderItem != null && productOrderItem.getId() != null) {
            log.info("productOrderItem is not  null, to update");
            updateProductOrderItem(taskInstance);
            return;
        }

        String host = taskInstance.getHost();
        if (StringUtils.isBlank(host) || host.indexOf(":") == -1) {
            log.error("create product order item error, host is null");
            return;
        }

        log.info("============= create product order item (" + taskInstance.getId() + ") start =============");
        log.info(JSONUtils.toJsonString(taskInstance));
        ProductOrder productOrder = productOrderMapper
                .selectProductOrderByProcessInstanceId(taskInstance.getProcessInstanceId(), taskInstance.getRunTimes());

        if (productOrder != null) {
            Date now = new Date();
            productOrderItem = new ProductOrderItem();
            productOrderItem.setTaskInstanceId(taskInstance.getId());
            productOrderItem.setNo(taskInstance.getTaskCode() + "");
            productOrderItem.setAmount(new BigDecimal(0));
            if (StringUtils.isNotBlank(host) && host.indexOf(":") != -1) {
                host = host.substring(0, host.indexOf(":"));
            }
            if (StringUtils.isNotBlank(host)) {
                Product product = productMapper.selectProductByHost(host);
                if (product != null) {
                    productOrderItem.setResourceNo(product.getProductNo()); // 根据hosts获取资源编号
                    productOrderItem.setProductName(product.getProductName());
                    productOrderItem.setHost(product.getHost());
                    productOrderItem.setProductType(product.getProductType());
                    productOrderItem.setProductSource(product.getProductSource());
                    productOrderItem.setBillType(product.getBillType());
                    productOrderItem.setBillRule(product.getBillRule());
                    productOrderItem.setPrice(product.getPrice());
                    productOrderItem.setvCpu(product.getvCpu());
                    productOrderItem.setMemory(product.getMemory());
                    productOrderItem.setDataTotal(product.getDataTotal());
                    productOrderItem.setServiceRatio(product.getServiceRatio());
                } else {
                    log.info("create product order item error, not found resource no");
                    return;
                }
            }
            productOrderItem.setStartTime(taskInstance.getStartTime());
            productOrderItem.setEndTime(taskInstance.getEndTime());
            productOrderItem.setTaskStatus(0); // '运行状态 0=运行中,1=运行完成,2=运行失败',
            productOrderItem.setProductOrderId(productOrder.getId());
            productOrderItem.setCreateTime(now);
            productOrderItem.setUpdateTime(now);
            productOrderItemMapper.insert(productOrderItem);
        } else {
            log.error("create product order item error, productOrder is null");
        }
        log.info("============= create product order item (" + taskInstance.getId() + ") success =============");
    }

    private void calcAmount(TaskInstance taskInstance, ProductOrderItem productOrderItem) {
        BigDecimal amount = new BigDecimal(0);
        BigDecimal initAmount = new BigDecimal(0);
        BigDecimal caclAmount = new BigDecimal(0);
        // 1.获取host,根据host获取资源详情
        String host = taskInstance.getHost();
        if (host.indexOf(":") != -1) {
            host = host.substring(0, host.indexOf(":"));
        }
        Product product = productMapper.selectProductByHost(host);
        if (product != null) {
            // 2.计算任务金额
            // 2.1计算任务执行时间，单位秒
            long taskTime = DateUtils.differSec(taskInstance.getEndTime(), taskInstance.getStartTime());
            // 2.2将任务执行时间转换为分钟 不足一分钟按一分钟计算 超过30S四舍五入
            long min = 0;
            if (taskTime >= 0 && taskTime <= 60) {
                min = 1;
            } else if (taskTime > 60) {
                min = taskTime / 60 + (taskTime % 60 >= 30 ? 1 : 0);
            }
            if (min > 0) {
                // 2.3计算任务金额=单价*时长
                initAmount = product.getPrice().multiply(new BigDecimal(min));
                amount = product.getPrice().multiply(new BigDecimal(min));
                // 获取任务计算规则
                QueryWrapper<ChargingRule> wrapper = new QueryWrapper<>();
                wrapper.eq("rule_type", product.getBillRule());
                wrapper.orderByDesc("min_number");
                List<ChargingRule> rules = chargingRuleMapper.selectList(wrapper);
                // log.info(product.getBillRule()+" rule size:"+rules.size());
                if (rules != null && rules.size() > 0) {
                    // 2.4计算计费规则，如果存在计费规则
                    BigDecimal max = new BigDecimal(100);
                    if (product.getBillRule().intValue() == 0) {
                        // 2.4.1 标准计费
                        ChargingRule chargingRule = rules.get(0);
                        if (chargingRule != null && chargingRule.getRebateRatio() != null) {
                            amount = amount.multiply(chargingRule.getRebateRatio()).divide(max,
                                    BigDecimal.ROUND_HALF_UP, 2);
                        }
                    } else {
                        // 2.4.2 梯度计费
                        for (ChargingRule chargingRule : rules) {
                            if (amount.compareTo(chargingRule.getMinNumber()) > 0) {
                                // 当前梯度内金额
                                BigDecimal amountLevel = null;
                                if (chargingRule.getMaxNumber().intValue() == -1) {
                                    // 最大值为无限大
                                    amountLevel = amount.subtract(chargingRule.getMinNumber());
                                } else {
                                    if (amount.compareTo(chargingRule.getMaxNumber()) > 0) {
                                        // 不在范围内的不计算
                                        caclAmount = caclAmount.add(amount.subtract(chargingRule.getMaxNumber()));
                                        amount = chargingRule.getMaxNumber();
                                    }
                                    amountLevel = chargingRule.getMaxNumber().subtract(chargingRule.getMinNumber());

                                }
                                // 计算折扣
                                caclAmount = caclAmount.add(amountLevel.multiply(chargingRule.getRebateRatio())
                                        .divide(max, BigDecimal.ROUND_HALF_UP, 2));
                                // 获取本次计算后剩余金额
                                amount = amount.subtract(amountLevel);
                            }
                        }
                    }
                }
            }
        }
        productOrderItem.setAmount(initAmount);
        productOrderItem.setRealAmount(amount.add(caclAmount));
        productOrderItem.setDiscountAmount(productOrderItem.getAmount().subtract(productOrderItem.getRealAmount()));
    }

}
